#!/usr/bin/python

from TimeChessUtil import *
from numpy import *
from pieces import *
import sys
from copy import *
import time

from UserArray import UserArray


#############################################################
#					TIMECHESSENGINE OBJECT					#
#############################################################
class TimeChessEngine(object):
	def __init__(self):

		# The main game array
		self.game = newGame()

		# The current turn
		self.t = 0

		# Which colour has the current turn
		self.currentSide = ""
		self.updateCurrentSide()

		# A place to store moves into the future until they occur
		self.futureQueue = []

		# Are we a copy that should not try to figure out checkmate?
		self.isCopy = False


	def getPossibleMoves(self, pos):
		p = self.game.getPieceAt(pos)
		if (p is None) or (p.pos.t != self.t):
			return []
		else:
			return p.getMoves()

	def attemptMove(self, move, ignoreCheck=False):
		debug("attempting move:\t%s" % move)
		
		success = self.checkMove(move)
		if success:
			self.__makeMove(move)
		return success

	def moveEndsInCheck(self, move):
		side = self.game[move.start].color
		testEngine = deepcopy(self)
		testEngine.isCopy = True
		debug("======================COPY __makemove STARTS======================")
		testEngine.__makeMove(move)
		debug("======================COPY __makemove ENDS ======================")
		debug("======================COPY __isInCheck STARTS ======================")
		if testEngine.isInCheck(side):
			debug("======================COPY __isInCheck ENDS ======================")
			debug("Move failed: ends with you in check")
			return True
		else:
			debug("======================COPY __isInCheck ENDS ======================")
			return False


	def checkMove(self,move):
		def fail(message):
			userMsg("ILLEGAL: move %s" % move)
			userMsg("\t%s"%message)
			return Result(False, message)
		debug("checking move legality:\t%s" % move)


		if not move.start.isOnBoard():
			return fail("%s is not a location on the board" % move)

		if move.start.t != self.t:
			return fail("It is turn %s, you cannot move a piece from turn %s"%
						(self.t,move.start.t))

		piece = self.game[move.start]
		if piece is None:
			return fail("There is no piece at %s" % move.start)

		if piece.color != self.currentSide:
			return fail("Piece at %s is %s, but it is %s's turn"%
							(move.start, piece.humanString(), self.currentSide))

		if move not in piece.getMoves():
			return fail("Not a possible move for %s at %s"%(piece.humanString(),piece.pos))
		target = self.game[move.end]
		if target and target.color == piece.color:
			return fail("Cannot take a piece which is on your side")

		if type(target) == King:
			return fail("You cannot take the %s at %s because it's a King"%
										(target.humanString(), move.end))

		if self.moveEndsInCheck(move):
			return fail("Move ends with you in check")

		debug("checkMove: move seems legal:\t%s" % move)
		return Result(True,"Success")
		
	def updateCurrentSide(self):
		if self.t%2:
			self.currentSide = "b"
		else:
			self.currentSide = "w"

	def newTurn(self):
		"""
		Tick over to a new turn, making a copy of the top layer and incrementing t
		"""

		# make a child for every piece in the current turn and put
		# the child in the new next turn
		for item in self.game[self.t].flatten():
			if item is not None:
				child = item.makeChild()
				if child: self.game[child.pos] = child


		# Now it's next turn
		self.t += 1
		self.updateCurrentSide()

		if (not self.isCopy) and self.isInCheckmate(self.currentSide):
			userMsg("%s is checkmated"%self.currentSide)
			self.gameOver(loser=self.currentSide)

		debug("TURN %s BEGINS"%self.t)


	def eraseLineageFrom(self, rootPiece):
		"""
		Remove all pieces that descend from the given piece

		If you call this function with a piece in the past, all its occurrences
		from then until the present are removed from the board.
		"""

		def isDescendant(piece):
			"""Is this piece a descendant of the root piece"""
			# Does it have the same pieceID and a greater age
			if piece:
				return ((piece.pieceId == rootPiece.pieceId) and
							(piece.age >= rootPiece.age))
			else:
				return False

		# make a mask of the `Game` with `True` over descendants
		mask = vectorize(isDescendant)(self.game)

#		#Iterate over the descendants
#		for t, y, x in zip(*nonzero(mask)):
#			print Pos(x, y, t)

		#Set all the non-false mask locations in the game to None, erasing the pieces
		self.game[nonzero(mask)] = None


	def doFutureMoves(self):
		"""
		Cause any time travelling pieces from the past due to arrive this turn to arrive
		"""
		self.ensureFutureMovesAreLegal()

		for p in self.futureQueue:
			if p.pos.t == self.t:
				debug("%s moves in from the past to %s" % (p.humanString(), p.pos))
				self.game[p.pos] = p
				self.futureQueue.remove(p)
		
	def ensureFutureMovesAreLegal(self):
		"""
		Make sure all the pieces queued in the futureQueue are still making legal moves

		A Future Move can become illegal after it has been made.
		If a piece tries to move into the future, and while it is 'in transit'
		an opposing piece moves in the way, the future move is illegal, and the
		future moving piece becomes "Lost in Time".
		"""
		for p in self.futureQueue:

			parent = self.game.getPieceParent(p)
			if (parent is None):
				userMsg("%s was destroyed in the past. Forward move to %s never happened" %
																(p.humanString(), p.pos))
				self.futureQueue.remove(p)

			elif p.pos not in [m.end for m in parent.getMoves()]:
				userMsg("%s's forward move to %s has become illegal, so it is now Lost in Time" %
																		(p.humanString(), p.pos))
				self.futureQueue.remove(p)
				if (type(p) == King) and (not self.isCopy):
					userMsg("%s has become Lost in Time!"%p.humanString())
					self.gameOver(loser=p.color)

			# if the move is a pawn
			elif (type(p)==Pawn):

				# if it's moving diagonally
				if (p.pos.x != parent.pos.x):
					target = self.game[Pos(p.pos.x,p.pos.y,p.pos.t-1)]
					if (target is None) or (target.color==p.color):
						userMsg("%s was moving diagonally into the future but there was nothing there to attack. The diagonal move was thus illegal, and the piece has become Lost in Time!"%p.humanString())
						self.futureQueue.remove(p)


	def isInCheckmate(self, color):
		if not self.isInCheck(self.currentSide):
			return False
		userMsg("Checking for Checkmate")
		possibleMovers = self.game.getPosWhereTrue((lambda x:(x.color==color) and (x.pos.t ==self.t)))

		for piece in [self.game[pos] for pos in possibleMovers]:
			for move in piece.getMoves():
				# No future move by anything other than a king can escape check
				if (type(piece) == King) or (move.end.t <= self.t):
					if not self.moveEndsInCheck(move):
						userMsg("Not checkmated, %s is legal"%move)
						return False
		
		# If we got through all the possible moves without finding one that didnt end in check
		# we are checkmated
		return True

	def isInCheck(self, color):
		"""
		Return True if the given colour is in check right now
		"""

		#make a lambda function that is given a piece and returns true if it's a potential threat
		isThreat = (lambda x: (x.pos.t==self.t) and (x.color != color))
#		threatenedPosSet = set([])

		# find all our king's positions, now and past
		isOurKing = (lambda x: (type(x)==King) and (x.color == color))
		kingPosSet = set(self.game.getPosWhereTrue(isOurKing))

		inCheck = False

		for ThreatPos in self.game.getPosWhereTrue(isThreat):
			threat = self.game[ThreatPos]
#			print "isincheck:", threat, threat.game is self.game, id(threat.game), id(self.game)
			for m in threat.getMoves():
#				threatenedPosSet.update([m.end])
				if m.end in kingPosSet:
					inCheck = True
					if not self.isCopy:
						userMsg("In check: %s: %s king at %s threatened by %s at %s"%(time.time(), color,m.end,threat.humanString(),threat.pos))
					else:
						sys.stderr.write(".")
					debug(str([(str(m.end),str(self.game[m.end])) for m in threat.getMoves()]))
					debug(str([(str(m.end),str(threat.game[m.end])) for m in threat.getMoves()]))
#					print id(threat.game)

		return inCheck
		

	def __makeMove(self, move):
		"""
		Make a move in the game.

		Note:	This method DOES NO CHECKING AT ALL. It will allow any move, and may
				crash with OutOfBoundsExceptions etc. IT SHOULD NEVER BE CALLED DIRECTLY.
				Use attemptMove(move) which does checks and then calls this
		"""
		if move.isPresent:		#if the move is a normal chess move
			debug("making move in the present: %s" % move)

			# Grab the piece we want to move
			#piece = self.game.getPieceAt(move.start)
			debug(str(type(self.game)))
			piece = self.game[move.start]

			# The place we moved to is now the piece
			self.game[move.end] = piece

			# The piece's position is its new one
			piece.pos = move.end

			# The place we moved from is now empty
			self.game[move.start] = None

			debug("Present move made, game state is:\n%s" % self.game.prettyPrint())

		# if this is a move into the past	
		elif move.isPast:
			debug("making move in the past: %s" % move)

			piece = self.game[move.start]
			self.game[move.start] = None

			dest = self.game[move.end]

			if dest is not None:	# If piece landed on something
				if dest.color != piece.color: # of opposite color
					debug("%s moves into the past and takes %s in %s" %
								(piece.humanString(), dest.humanString(), move.end))
					self.eraseLineageFrom(dest) # wipe it out from here on
				else:
					raise Exception("Tried to take a piece of your own colour in the past")
			
			# Propagate the piece to the present as a braindead piece
			x = move.end.x
			y = move.end.y
			t = move.end.t
			curPiece = piece
			for t in range(t, self.t+1): #from the time of capture until the present
				pos = Pos(x,y,t)
				if self.game[pos] is None:
					curPiece = curPiece.makeChild()
					curPiece.pos = pos
					self.game[pos] = curPiece
				else:
					break

			debug("Past move made, game state is:\n%s" % self.game.prettyPrint())

		elif move.isFuture:
			debug("making move in the future: %s" % move)

			piece = self.game[move.start]

			#don't get rid of it, sterilise it
#			self.game[move.start] = None

			newpiece = piece.makeChild()
			piece.sterile = True

			newpiece.pos = move.end

			self.futureQueue.append(newpiece)

			debug("current future queue:")
			for p in self.futureQueue:
				debug("\t%s to appear at %s " % (p.humanString(), p.pos))


		self.doFutureMoves()


		self.newTurn()

	def gameOver(self, winner=None, loser=None):
		if (winner is None) and (loser is None):
			userMsg("Game Over. No Winner")
		elif (winner == "b") or (loser == "w"):
			userMsg("Game over, Black Wins")
		elif (winner == "w") or (loser == "b"):
			userMsg("Game over, White Wins")
		else:
			userMsg("Game Over.")

		self.reset()
	
	def reset(self):
		userMsg("Starting New Game")
		self.__init__()

#############################################################
#						GAME OBJECT							#
#############################################################

class Game(UserArray):
	def __init__(self, *args, **kwargs):
		"""
		Create the game array.
		Do not put actual initialisation in here, it gets called all the time by the
		numpy system
		"""
#		debug("creating new Game object")
		super(Game, self).__init__(*args, dtype=object, **kwargs)
	
	def getPieceAt(self, pos):
		try:
			return self[pos]
		except IndexError:
			return None

	def __getitem__(self, index):
		"""Override the indexing to handle being given a Pos object"""
		if type(index) == Pos:
			try:
				return super(Game,self).__getitem__(index.t)[index.y][index.x]
			except IndexError:
				return None
		else:
			return super(Game,self).__getitem__(index)

	def __setitem__(self, index, value):
		"""Override the indexing to handle being given a Pos object"""
		if type(index) == Pos:
			self.array[index.t,index.y,index.x] = value
		else:
			self.array[index] = value
	
	def getBoardAtTurn(self, turn):
		return Game([self[turn]])

	
	def prettyPrint(self,selected = None):
		t = 0
		string = ""
		while sometrue(self[t].flatten()):
#			string +=	"\nTurn %s:\n" % t
#			string +=	"    a b c d e f g h \n"
			string +=	"    _______________ ______Turn %s\n" % t
			for rank in "87654321":
				string += "%s  " % rank
				for file in "abcdefgh":
					pos = Pos(file+rank+"t"+str(t))
					startcode = ""
					if selected:
						if (pos == selected.pos):
							startcode += "\033[0;44m"
						if pos in selected.getMoves():
							startcode += "\033[0;42m"
							if self[pos] and self[pos].color != selected.color:
								startcode += "\033[0;41m"
					endcode = "\033[0;0m"
					string += "|%s%s%s" % (startcode,self[pos] or "_", endcode)
				string += "|  %s\n" % rank
			string +=	"\n    a b c d e f g h\n"

			t += 1

		return string

	def getPosWhereTrue(self,f):
		"""
		Return a list of positions at which the function f holds true
		"""
		#ensure the value isn't none, so we can get attributes of Piece without error
		nf = (lambda i: (i is not None) and f(i))
		vf = vectorize(nf)
		indices = nonzero(vf(self))

		# Roll the coordinate arrays into a list of Pos objects
		return [Pos(p[2],p[1],p[0]) for p in zip(*indices)]

	def getPieceParent(self, p):
		"""
		Return the position of the direct parent of piece p (the instance one turn before p)
		"""
		debug("getPieceParent called on %s of age %s"%(p.humanString(), p.age))
		isParent = (lambda x:(x.pieceId==p.pieceId) and (x.age==(p.age-1)))
		poslist = self.getPosWhereTrue(isParent)
		debug("poslist: %s"%poslist)
		if poslist and poslist[0]:
			return self[poslist[0]]
		else:
			return None


	def __deepcopy__(self,memo):
		"""
		Return a full copy of the Game

		Used by the copy module to copy the object and all sub-objects
		"""
		newarr = copy(self)

		def makeCopy(p):
			"Small function to make a copy of a Piece object"
			if p:
				return type(p)(p.color,p.pos,self,p.pieceId,p.age)
			else:
				return None
		f = vectorize(makeCopy)
		newGame = Game(f(newarr))

		# Set the game variable for every piece to the game they're in
		# since they will be pointing to the old game object at first
		indices = nonzero(newGame)
		for pos in [Pos(p[2],p[1],p[0]) for p in zip(*indices)]:
			newGame[pos].game = newGame

		return newGame




# MODULE LEVEL FUNCTIONS
def newGame():
	game = Game(zeros((500,8,8), object))

	game[:,:,:] = None

	for y, col in [(1,"w"),(6, "b")]:
		for x in range(8):
			game[0,y,x] = Pawn(col,Pos((x,y,0)),game)

	for y, col in [(0,"w"),(7, "b")]:
		for x, constructor in enumerate([Rook,Knight,Bishop,Queen,King,Bishop,Knight,Rook]):
			game[0,y,x] = constructor(col,Pos((x,y,0)),game)

	for y in [2,3,4,5]:
		for x in range(8):
			game[0,y,x] = None

	

	debug("New Game made:\n%s" % game.prettyPrint())
	return game




#====================================================================
#                    UNIT TESTS
#====================================================================
import unittest

class TestGame(unittest.TestCase):

	def setUp(self):
		self.g = newGame()

	def test_getPieceAt(self):
		
		# Asking for completely absurd positions returns None
		self.assertEqual(self.g.getPieceAt(Pos("i0t9")), None)

		# a Pawn here
		self.assertEqual(type(self.g.getPieceAt(Pos("a2t0"))), Pawn)

		# nothing here
		self.assertEqual(self.g.getPieceAt(Pos("a3t0")), None)

	def test_getBoardAtTurn(self):

		# always return a Game object
		self.assertEqual(type(self.g.getBoardAtTurn(0)), Game)
		self.assertEqual(type(self.g.getBoardAtTurn(1)), Game)

		# even when the turn number is stupid
		self.assertEqual(type(self.g.getBoardAtTurn(-1)), Game)

		#the game object has only one layer
		self.assertEqual(len(self.g.getBoardAtTurn(0)), 1)
	
	def test_prettyprint(self):
		self.assertEqual()



	def test_newGame(self):
		g = newGame()

		# Should make a Game not an Array
		self.assertEqual(type(g), Game)


class TestTCE(unittest.TestCase):

	def setUp(self):
		self.e = TimeChessEngine()
	
	def tearDown(self):
		self.e = None

	def test_init(self):

		self.assertEqual(self.e.currentSide, "w")
	
	def test_gpm(self):

		# equivalent ways of getting moves w/ good pos
		p = Pos("a2t0")
		self.assertEqual(self.e.getPossibleMoves(p), self.e.game[p].getMoves())

		# Bad pos is an empty list
		p = Pos("a2t30")
		self.assertEqual(self.e.getPossibleMoves(p), [])

	def test_attemptMove(self):
		t = self.e.t
#		self.e.attemptMove(Move(Pos(""),Pos("")))

		# legal move
		self.e.attemptMove(Move(Pos("a2t0"),Pos("a4t0")))

		# a turn has passed
		self.assertEqual(self.e.t, t+1)

		# the piece is not where it was
		self.assertEqual(self.e.game[Pos("a2t0")], None)

		# there is a piece in the new position
		self.assertTrue(issubclass(type(self.e.game[Pos("a4t0")]), Piece ))


		# check the age rules
		self.assertEqual(self.e.game[Pos("a4t0")].age, 0)
		self.assertEqual(self.e.game[Pos("a4t1")].age, 1)

		# they have the same ID
		self.assertEqual(self.e.game[Pos("a4t0")].pieceId, self.e.game[Pos("a4t1")].pieceId)

	def test_checkMove(self):
		
#		self.assertFalse(self.e.checkMove(Move(Pos(""),Pos("")))
		#start isn't on the board, move will fail
		self.assertFalse(self.e.checkMove(Move(Pos("e9t0"),Pos("a2t4"))))

		# not the right turn
		self.assertFalse(self.e.checkMove(Move(Pos("a2t3"),Pos("a3t3"))))

		# no piece to move
		self.assertFalse(self.e.checkMove(Move(Pos("a4t0"),Pos("a5t0"))))

		# moving the wrong side's pieces
		self.assertFalse(self.e.checkMove(Move(Pos("a7t0"),Pos("a6t0"))))

		# making a move that piece can't amke
		self.assertFalse(self.e.checkMove(Move(Pos("a2t0"),Pos("b4t0"))))

		# try to take your own piece
		self.assertFalse(self.e.checkMove(Move(Pos("b1t0"),Pos("d2t0"))))

	def test_newTurn(self):
		# manually increment the turn
		self.e.newTurn()

		self.assertEqual(self.e.t, 1)
		self.assertEqual(self.e.currentSide, "b")

		self.assertEqual(self.e.game[Pos("a2t1")].age, 1)
	
	def test_eraseLineage(self):

		# ff 5 turns
		for i in range(5):
			self.e.newTurn()

		# put a bishop where it can attack
		self.e.game[Pos("c7t5")] = Bishop("b",Pos("c7t5"), self.e.game)

		pid = self.e.game[Pos("c2t5")].pieceId
		self.e.attemptMove(Move(Pos("c7t5"),Pos("c2t0")))

		# piece is totally wiped out
		self.assertEqual(self.e.game.getPosWhereTrue((lambda x:x.pieceId == pid)),
								[])

	def test_future(self):
		# clear a gap for the queen
		self.e.game[Pos("d2t0")] = None

		q = self.e.game[Pos("d1t0")]

		# move queen to the future
		self.e.attemptMove(Move(Pos("d1t0"),Pos("d4t3")))

		# queue is now 1 long
		self.assertTrue(self.e.futureQueue)
		self.assertEqual(len(self.e.futureQueue), 1)

		# same pieceId
		self.assertEqual(self.e.futureQueue[0].pieceId, q.pieceId)

		# age +1
		self.assertEqual(self.e.futureQueue[0].age, q.age+1)

		# not there yet
		self.assertFalse(self.e.game[Pos("d4t3")])
		
		while self.e.t != 3:
			self.e.newTurn()

		self.e.doFutureMoves()

		# now the piece arrives
		self.assertEqual(self.e.game[Pos("d4t3")].pieceId, q.pieceId)
	
	def test_lost_in_time(self):
		q = self.e.game[Pos("d1t0")]

		# move queen to the future
		self.e.attemptMove(Move(Pos("d1t0"),Pos("d4t3")))

		# queue is now 1 long
		self.assertTrue(self.e.futureQueue)
		self.assertEqual(len(self.e.futureQueue), 1)

		# not there yet
		self.assertFalse(self.e.game[Pos("d4t3")])
		
		while self.e.t != 3:
			self.e.newTurn()

		# piece gets lost in time
		self.e.doFutureMoves()

		# Piece never arrives
		self.assertEqual(self.e.game[Pos("d4t3")], None)


	def test_isincheck(self):
		# wipe the game
		self.e.game[:,:,:] = None

		# set up a check position
		self.e.game[Pos("c7t0")] = King("b", Pos("c7t0"), self.e.game)
		self.e.game[Pos("c2t0")] = King("w", Pos("c2t0"), self.e.game)
		self.e.game[Pos("e7t0")] = Rook("w", Pos("e7t0"), self.e.game)


		# black is in check
		self.assertTrue(self.e.isInCheck("b"))

		# white is not
		self.assertFalse(self.e.isInCheck("w"))

	def test_isincheckmate(self):

		self.e.game[Pos("e7t0")] = Bishop("b",Pos("e7t0"), self.e.game)
		self.e.game[Pos("e2t0")] = None

		for i in range(6):
			self.e.newTurn()

		# turn is now 0 because checkmate reset gaem
		self.assertEqual(self.e.t, 0)




		

if __name__ == '__main__':
	unittest.main()

